import torch
from torch import nn
from src.model.space_model.gconv import Gconv
from typing import Tuple


class Gcn(nn.Module):
    def __init__(self, dims: Tuple[int, int, int]):
        super(Gcn, self).__init__()
        self.input_size, self.hidden_size, self.output_size = dims

        self.gc = Gconv(dims)
        self.dropout = nn.Dropout(0.1)  # 添加Dropout层
        
    def forward(self, x, adj):
        # x的尺寸是(B, P, N, H)
        out = self.gc(x, adj)
        # 最终的out是(batch, N, hidden)
        return out
